# cython: profile=True
from v3dsolver.intvector cimport IntVec,IntVec_new
from v3dsolver.vector cimport Vec
from v3dsolver.fluid_elements.fluid_elems cimport Saffman,FluidElem
from mpi4py cimport MPI
from numpy cimport ndarray
cimport cython

cdef extern from "math.h":
    double sqrt(double)

cdef class Bin:
    cdef public list elems
    cdef public list domain
    cdef public MPI.Comm comm
    cdef public double cellsize
    cdef public list intdomain
    cdef public ncells_x
    cdef public ncells_y
    cdef public ncells_z
    cdef public list hashelems
    cdef public size
    cdef public dict first_elem_indices
    cdef public dict next_elem_indices
    cdef public dict elem_count
    cdef public dict nbrlist
    

    cdef public list assign_nbrs(self)
                                 
    @cython.locals(xmin=double,ymin=double,zmin=double)
    cpdef IntVec get_index(self,Vec pos)
    
    @cython.locals(elem_x=double,elem_y=double,elem_z=double)
    cdef Vec get_newelem_pos(self,index_1d)
    

cdef class NeighbourFinder:
    cdef public Bin bin
    cdef public list elems
    cdef public list hashelems
    cdef public double R
    cdef public double nu
    cdef public double dt
    cdef public double cellsize
    cdef public pop_data
    cdef public list nbr_elems
    cdef public ncells_x
    cdef public ncells_y
    cdef public ncells_z
    cdef public dict cellelems_list
    cdef public dict remote_nbrs
    
    #@cython.locals(pos=Vec,parentcell=IntVec,index_1d=int,nbrs=list,
    #               nbr_cell=int,i=int)
    #cdef get_neighbours_serial(self,Saffman element)
    
    @cython.locals(proc_map=list,i=int,ith_proc_map=dict,
                   j=int,nbr=int,send_data=dict)
    cdef send_particles(self, MPI.Comm comm)
    
    @cython.locals(cellelems=list,current_elem=Saffman)
    cdef _update_elems(self,index_1d)
    
    @cython.locals(rvec=Vec,r=double)
    cdef bint are_neighbours(NeighbourFinder self, FluidElem base_elem, dict hashnbr)
